#pragma once
#include "copy.h"
#include <iostream>

/** 
 * Wrapper around cublasSgbmv().
 * Performs Z = alpha*A*X + beta*Y.
 *
 * @param _transA: whether or not to transpose matrix A
 * @param m: number of rows of the matrix A
 * @param ku: number of super-diagonals
 * @param alpha: scaling factor for A
 * @param A: matrix supplied in Diagonal format
 * @param X: vector to be multiplied by A
 * @param beta: scaling factor for Y
 * @param Y: vector to be added to A*X
 * @param Z: preallocated buffer for result
 */
void sgbmv
(
    int _transA,
    int m,
    int ku,
    float alpha,
    Matrix A,
    Vector X,
    float beta,
    Vector Y,
    Vector Z
)
{
    int n, kl, lda, incx, incy, size_y;
    char transA;

    if(A.stride != 1){
        //warn user
        std::cerr << "sgbmv: WARNING: Passing in a row major matrix has incurred a copy operation" << std::endl;
        int size_A = A.sub.length * A.length * sizeof(float);
        float *A_data;
        cudaMalloc(&A_data, size_A);
        SubMatrix A_sub(A.sub.length, A.length, A_data, 0);
        Matrix A_copy(A.length, 1, A_sub);
        copy(A, A_copy);
        A = A_copy; //huzzah for args passed by copy
    }

    n = A.sub.length;
    kl = A.length - ku - 1;
    lda = A.sub.stride;
    incx = X.stride;
    incy = Y.stride;

    if(_transA){
        if(m != X.length){
            std::cerr << "sgbmv: ERROR: mismatch in demensions of A and X (" << m
                << " and " << X.length << " should be equal)" << std::endl;
        }

        if(n != Y.length){
            std::cerr << "sgbmv: ERROR: mismatch in demensions of A and Y (" << n 
                << " and " << Y.length << " should be equal)" << std::endl;
        }
    } else {
        if(n != X.length){
            std::cerr << "sgbmv: ERROR: mismatch in demensions of A and X (" << n
                << " and " << X.length << " should be equal)" << std::endl;
        }

        if(m != Y.length){
            std::cerr << "sgbmv: ERROR: mismatch in demensions of A and Y (" << m 
                << " and " << Y.length << " should be equal)" << std::endl;
        }
    }

    size_y = (1 + (Y.length * incy)) * sizeof(float);
    cudaMemcpy(Z.data, Y.data, size_y, cudaMemcpyDeviceToDevice);

    transA = _transA ? 't' : 'n';

    cublasSgbmv(transA, m, n, kl, ku, alpha, A.sub.data, lda, X.data, incx, beta, Z.data, incy);
}

void sgbmv
(
    int _transA,
    int m,
    int ku,
    float alpha,
    Matrix A,
    stored_sequence<float> X,
    float beta,
    stored_sequence<float> Y,
    stored_sequence<float> Z
)
{
    Vector x(X.length, 1, X.data);
    Vector y(Y.length, 1, Y.data);
    Vector z(Z.length, 1, Z.data);
    sgbmv(_transA, m, ku, alpha, A, x, beta, y, z);
}
